########################################
## main_simulations.R
## File to run the main simulations from
## Blackwell and Olson (2021)
########################################

library(glmnet)
library(inters)
library(sandwich)
library(KRLS)
library(BART)

args <- commandArgs(TRUE)
if (length(args) == 0) args <- c(0, 425, 0, 0, 20, 100, "main")
mc.cores <- 1
source("getcoef_sim.R")

run <- as.integer(args[1])
n_units <- as.integer(args[2])
R22  <- as.numeric(args[3]) ## R^2 for Y
R21  <- as.numeric(args[4]) ## R^2 for D
n_covs  <- as.integer(args[5])
sims <- as.integer(args[6])
type <- args[7]
n_boots <- 500

## original simulations run on R 3.5.1. After 3.6.0, R changed the RNG
## for the sample function, which leads to non-reproducible results
## for the simulations.
if (getRversion() >= "3.6.0") RNGkind(sample.kind = "Rounding")
set.seed(12345 + run)

if (!dir.exists("output")) dir.create("output")
out_file <- paste0(
  "output/", type, "_sim_", n_units, "_",
  R22, "_", R21, "_", n_covs, ".csv"
)

coef_lasso <- matrix(NA, nrow = sims, ncol = 2)
coef_alasso <- matrix(NA, nrow = sims, ncol = 2)
coef_bart <- matrix(NA, nrow = sims, ncol = 2)
coef_pds <- matrix(NA, nrow = sims, ncol = 2)
coef_ols <- matrix(NA, nrow = sims, ncol = 2)
coef_sing <- matrix(NA, nrow = sims, ncol = 2)
coef_oracle <- matrix(NA, nrow = sims, ncol = 2)
ses_lasso <- matrix(NA, nrow = sims, ncol = 2)
ses_lasso <- matrix(NA, nrow = sims, ncol = 2)
ses_pds <- matrix(NA, nrow = sims, ncol = 2)
ses_ols <- matrix(NA, nrow = sims, ncol = 2)
ses_sing <- matrix(NA, nrow = sims, ncol = 2)
ses_oracle <- matrix(NA, nrow = sims, ncol = 2)
alasso_cov <- matrix(NA, nrow = sims, ncol = 2)
bart_cov <- matrix(NA, nrow = sims, ncol = 2)

# Covariates X are drawn from a multivariate normal distribution
sigma <- toeplitz(0.5 ^ (0:(n_covs - 1)))

betas <- get_coefs(type = type, n_covs,
                   R2_d = R21, R2_y = R22)

beta_y  <- betas$beta_y
beta_v  <- betas$beta_v
beta_d  <- betas$beta_d
beta_vy <- betas$beta_vy
beta_dy <- betas$beta_dy
beta_vd <- betas$beta_vd


# set up treatment effect and moderator effect
treatment_effect <- 0.5
moderator_effect <- 0.25
alpha0 <- 1
set.seed(02138 + run)

sim_iter <- function(s, ...) {
  if ((s / sims * 100) %% 1 == 0) {
    cat("Iteration: ", s, "/", sims, "\n")
  }
  out <- data.frame(n_units = n_units, n_covs = n_covs, dgp = type,
                    R2_d = R21, R2_y = R22)

  ## DGP
  X <- MASS::mvrnorm(n = n_units, mu = rep(0, n_covs), Sigma = sigma)
  X <- scale(X, scale = FALSE)
  c_names <- paste("X", 1:n_covs, sep = "")
  colnames(X) <- c_names
  # now generate V (moderator) and D (treatment) using X and relevant betas
  V <- rbinom(n_units, 1, boot::inv.logit(X %*% beta_v))
  XV <- X * V
  colnames(XV) <- paste("XV", 1:n_covs, sep = "")
  d <- X %*% beta_d + V * 0.25 + XV %*% beta_vd + rnorm(n_units)
  dX <- c(d) * X
  colnames(dX) <- paste("dX", 1:n_covs, sep = "")
  e_y <- X %*% beta_y + d * treatment_effect + V * moderator_effect +
    alpha0 * d * V + (XV) %*% beta_vy + dX %*% beta_dy
  y_sd <- 1
  y <- e_y  + rnorm(n_units, sd = y_sd)

  sim_data <- data.frame(y = y, d = d, d_V = c(d) * V, V = V, X, XV)

  ## oracle
  y_ora <- y - V * moderator_effect - X %*% beta_y -
    XV %*% beta_vy - dX %*% beta_dy
  oracle_out <- lm(y_ora ~ d * V)
  oracle_vcv <- sandwich::vcovHC(oracle_out)
  out$oracle_main <- oracle_out$coefficients["d"]
  out$oracle_int <- oracle_out$coefficients["d:V"]
  out$oracle_mainse <- sqrt(oracle_vcv["d", "d"])
  out$oracle_intse <- sqrt(oracle_vcv["d:V", "d:V"])


  ## post-double selection
  pds_out <- post_ds_interaction(data = sim_data, treat = "d", moderator = "V",
                                   outcome = "y", control_vars = c(c_names),
                                 method = "double selection")
  out$pds_main <- pds_out$coefficients["d"]
  out$pds_int <- pds_out$coefficients["d_V"]
  out$pds_mainse <- pds_out$se["d"]
  out$pds_intse <- pds_out$se["d_V"]

  ## design matrix for other lassos
  big_x <- cbind(d = c(d), d_V = c(d) * V, V = V, X, XV)
  small_x <- cbind(d = c(d), V = V, X = X)

  ## post-lasso
  lasso_out <- post_ds_interaction(
    data = sim_data, treat = "d", moderator = "V",
    outcome = "y", control_vars = c(c_names),
    method = "single selection")
  out$lasso_main <- lasso_out$coefficients["d"]
  out$lasso_int <- lasso_out$coefficients["d_V"]
  out$lasso_mainse <- lasso_out$se["d"]
  out$lasso_intse <- lasso_out$se["d_V"]


  ## fully moderated
  fully <- lm(y ~ ., data = sim_data)
  fully_vcv <- sandwich::vcovHC(fully)
  out$fully_main <- fully$coefficients["d"]
  out$fully_int <- fully$coefficients["d_V"]
  out$fully_mainse <- sqrt(fully_vcv["d", "d"])
  out$fully_intse <- sqrt(fully_vcv["d_V", "d_V"])

  ## krls
  krls_out <- KRLS::krls(y = y, X = small_x,
                               derivative = FALSE)
  small_x_pred <- rbind(cbind(d = 1, V = 1, small_x[, -c(1, 2)]),
                        cbind(d = 0, V = 1, small_x[, -c(1, 2)]),
                        cbind(d = 1, V = 0, small_x[, -c(1, 2)]),
                        cbind(d = 0, V = 0, small_x[, -c(1, 2)]))
  Xmeans <- colMeans(krls_out$X)
  Xsd <- apply(krls_out$X, 2, sd)
  XX <- scale(krls_out$X, center = Xmeans, scale = Xsd)
  small_x_pred <- scale(small_x_pred, center = Xmeans, scale = Xsd)
  nn <- nrow(small_x_pred)
  small_x_pred <- rbind(small_x_pred, XX)
  tmp <- matrix(0, nn + nrow(XX), nn + nrow(XX))
  tmp2 <- which(row(tmp) > col(tmp))
  tmp[tmp2] <- dist(small_x_pred)
  tmp <- tmp + t(tmp)
  tmp <- exp(-1 * tmp^2 / krls_out$sigma)

  newdataK <- matrix(tmp[1:nn, (nn + 1):(nn + nrow(XX))],
                     nrow = nn, byrow = FALSE)
  yfitted <- newdataK %*% krls_out$coeffs
  yfitted <- (yfitted * apply(krls_out$y, 2, sd)) + mean(krls_out$y)
  vcov.c.raw <-  krls_out$vcov.c * as.vector((1 / var(krls_out$y)))
  vcov.fit <- (apply(krls_out$y, 2, sd)^2) *
    tcrossprod(newdataK %*% vcov.c.raw, newdataK)
  ws <- c(1 / n_units, -1 / n_units, -1 / n_units, 1 / n_units)
  h <- matrix(rep(ws, each = n_units), ncol = 1)

  out$krls_int <- as.vector(t(h) %*% yfitted)
  out$krls_intse <- as.vector(sqrt(t(h) %*% vcov.fit %*% h)) * sqrt(4)
  out$krls_cov <- as.numeric(((out$krls_int - out$krls_intse * 1.96) < 1) &
                               ((out$krls_int + out$krls_intse * 1.96) > 1))


  ## adaptive lasso
  ridge_out <- glmnet::cv.glmnet(x = big_x, y = y, alpha = 0)
  ridge_coefs <- as.numeric(coef(ridge_out, s = "lambda.1se"))[-1]
  alasso_pens <- 1 / abs(ridge_coefs)
  alasso_pens[1:(ncol(X) + 3)] <- 0
  alasso_out <- glmnet::cv.glmnet(x = big_x, y = y, penalty = alasso_pens)
  alasso_coefs <- as.numeric(coef(alasso_out, s = "lambda.1se"))[-1]
  names(alasso_coefs) <- colnames(big_x)
  alasso_fits <- predict(alasso_out, newx = big_x, s = "lambda.1se")
  eps_alasso <- y - alasso_fits

  alasso_boots <- matrix(NA, nrow = n_boots, ncol = 2)
  for (b in seq_len(n_boots)) {
    star <- sample(seq_len(nrow(small_x)), replace = TRUE)
    y_star <- alasso_fits + eps_alasso[star]
    ridge_star <- glmnet::glmnet(x = big_x, y = y_star, alpha = 0,
                                 lambda = ridge_out$lambda.1se)
    ridge_coefs_star <- as.numeric(coef(ridge_star))[-1]
    alasso_pens <- 1 / abs(ridge_coefs_star)
    alasso_pens[1:(ncol(X) + 3)] <- 0
    alasso_star <- glmnet::glmnet(x = big_x, y = y_star, penalty = alasso_pens,
                                  lambda = alasso_out$lambda.1se)
    alasso_coefs_star <- as.numeric(coef(alasso_star))[-1]
    names(alasso_coefs_star) <- colnames(big_x)
    alasso_boots[b, ] <- alasso_coefs_star[c("d", "d_V")]
  }
  ci_alasso <- quantile(alasso_boots[, 2], probs = c(0.025, 0.975))

  out$alasso_main <- alasso_coefs["d"]
  out$alasso_int <- alasso_coefs["d_V"]
  out$alasso_intse <- sd(alasso_boots[, 2])
  out$alasso_cov <- as.numeric(ci_alasso[1] < 1 & ci_alasso[2] > 1)

  ## bart
  small_x_bart <- rbind(
    c(d = 1, V = 1, colMeans(X)),
    c(d = 0, V = 1, colMeans(X)),
    c(d = 1, V = 0, colMeans(X)),
    c(d = 0, V = 0, colMeans(X))
  )
  capture.output(
    bart_out <- BART::wbart(small_x, y, ndpost = 1000,
                            x.test = small_x_bart),
    file = "/dev/null"
  )
  bart_main <- bart_out$yhat.test[, 3] - bart_out$yhat.test[, 4]
  bart_int <- bart_out$yhat.test[, 1] - bart_out$yhat.test[, 2] -
    bart_out$yhat.test[, 3] + bart_out$yhat.test[, 4]
  out$bart_main <- mean(bart_main)
  out$bart_int <- mean(bart_int)
  bart_int_ci <- quantile(bart_int, probs = c(0.025, 0.975))
  out$bart_cov <- as.numeric(bart_int_ci[1] < 1 &
                                     bart_int_ci[2] > 1)

  ## single interaction
  single <- lm(y ~ d + d_V + V + X, data = sim_data)
  single_vcv <- sandwich::vcovHC(single)
  out$single_main <- single$coefficients["d"]
  out$single_int <- single$coefficients["d_V"]
  out$single_mainse <- sqrt(single_vcv["d", "d"])
  out$single_intse <- sqrt(single_vcv["d_V", "d_V"])

  return(out)
}

sim_out <- parallel::mclapply(1:sims, sim_iter, mc.cores = mc.cores)
out <- do.call(rbind, sim_out)
write.csv(out, file = out_file, row.names = FALSE)
